import torch
from torch.utils.data import DataLoader

from .worker import Worker
from ..summarizer import Summarizer


class Validator(Worker):
    def __init__(self, dataloader: DataLoader, summarizer: Summarizer) -> None:
        super().__init__(dataloader)
        self.summarizer = summarizer

    @torch.no_grad()
    def process_batch(self, batch_data):
        inputs, targets, labels = batch_data
        outputs, loss_dict = self.runner.model.forward(
            inputs, targets, use_amp=self.runner.enable_amp
        )
        self.summarizer.update(loss_dict, outputs, labels)

    def before_work_hook(self) -> None:
        self.runner.model.eval()
        return

    def after_work_hook(self) -> None:
        loss_summary, metrics_summary = self.summarizer.summarize()

        print("metrics")
        if metrics_summary:
            for k, v in metrics_summary.items():
                print(k, v)

        for sl in self.runner.sub_loggers:
            sl.log_val(loss_summary, self.runner.epoch_current, metrics_summary)
        return

    def get_desc(self) -> str:
        return self.DESC_TEMPLATE.format(job_name="  Val", step=self.runner.epoch_current)
